import numpy as np

"""
实现构造k-d树
"""


class KDTree(object):
    class Node(object):
        def __init__(self, point, fea_idx, instance):
            """
            创建节点
            :param point: 划分点
            :param fea_idx: 划分特征索引
            :param instance: 节点数据
            """
            self.point = point
            self.instance = instance
            self.fea_idx = fea_idx
            self.right = None
            self.left = None

        def __str__(self):
            return "该节点分割点为 %s, 分割特征索引为 %s， 节点实例为 %s" % (self.point, self.fea_idx, self.instance)

    def __init__(self):
        pass

    def build_tree(self, x):
        """
        创建树
        :param x:
        :return:
        """
        point = None
        fea_idx = None
        instance = None
        if len(x) == 1:
            return "完成"
        node = KDTree.Node(point, fea_idx, instance)
        # 1.得到方差最大的特征索引
        max_var_idx = self._max_var(x)
        node.fea_idx = max_var_idx
        # 2.获取该特征所有值，并获取其中位数和其索引
        best_feature = x[:, max_var_idx]
        point, idx = self._mid_fea(best_feature)
        node.point = point
        instance = x[idx]
        node.instance = instance
        idx_ = best_feature < point
        self.build_tree(x[idx_])
        self.build_tree(x[~idx_])
        return node

    @staticmethod
    def _mid_fea(feature):
        """
        获取该特征的中位数和其索引
        :param feature:
        :return:
        """
        le = len(feature)
        sort_fea = np.unique(feature)
        # 1.获取中位数（分割点）
        point = sort_fea[(le // 2)]
        # 获取该分割点的索引
        idx = np.where(feature == point)[0][0]
        return point, idx

    @staticmethod
    def _max_var(x):
        """
        计算最大方差，返回最大方差的特征索引
        :param x:
        :return:
        """
        var_ = np.apply_along_axis(np.var, axis=0, arr=x)
        return np.argmax(var_)


if __name__ == '__main__':
    da_1 = np.array([[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]])
    da_left_1 = np.array([[2, 3], [5, 4], [4, 7]])
    da_right_1 = np.array([[9, 6], [8, 1]])
    kd = KDTree()
    print(kd.build_tree(da_1))
    print(kd.build_tree(da_left_1))
    print(kd.build_tree(da_right_1))
